import os
import pymysql  
import random

from analysis_core import log_response
  

# 数据库连接配置  
db_config = {  
    'host': 'localhost',  
    'user': 'root',  
    'password': 'zhejianglab',  # 请确保以安全的方式处理密码  
    'database': 'tee_cipher',  
    'charset': 'utf8mb4',  
    'cursorclass': pymysql.cursors.DictCursor,  
}  
  
# 查询语句  
MYSQL_QUERY_LIST = [
    "select * from protein_demo where token in ('A0A4Y4LZJ8', 'A0A329ZC87');",
    "select * from protein_demo where name='A0A1F3WWV3_9PROT CRISPR-associated endonuclease Cas9';",
    "select * from protein_demo where sequence like '%EALDAKTKVSKNQARRTARAMRRQVARHKMRRNAALRYLTNNNLLPADQVE%';",
    "select * from protein_demo where attribution like '%OX=86473%'",
] 

# 执行结果文件存储位置, 使用os.path.expanduser来解析用户主目录，python并不能解析~字符为当前的用户主目录
MYSQL_RESULT_SAVE = os.path.expanduser("~/yemd/analysis-tee-scripts/result.csv")
print("sql result.csv path:", MYSQL_RESULT_SAVE)


def sql_execute_command(sql: str) -> int:
    affect_rows = 0
    # 打开文件以写入结果  
    with open(MYSQL_RESULT_SAVE, 'w') as output_file:  
        # 连接到数据库  
        connection = pymysql.connect(**db_config)  
        try:  
            with connection.cursor() as cursor:  
                # 执行查询  
                affect_rows = cursor.execute(sql) 
                print(f"sql affect_rows: %s" % affect_rows) 
                # 获取所有结果  
                result = cursor.fetchall()  
                # 先写表的title字段
                output_file.write(','.join([desc[0] for desc in cursor.description]) + '\n')
                # 将结果写入文件  
                for row in result: 
                    # print(repr(row))  
                    output_file.write(','.join(map(str, row.values())) + '\n')  
        finally:  
            # 关闭数据库连接  
            connection.close()
    return affect_rows


@log_response(key="SQL")
def random_sql():
    return random.choice(MYSQL_QUERY_LIST)


def main():
    sql_execute_command(random_sql())


if __name__ == "__main__":
    main()            